package com.cheems.cheemsspringaialibaba.Graph.node;

import com.alibaba.cloud.ai.graph.OverAllState;
import com.alibaba.cloud.ai.graph.action.NodeAction;
import org.springframework.ai.chat.client.ChatClient;

import java.util.HashMap;
import java.util.Map;

import static com.cheems.cheemsspringaialibaba.Graph.config.HenauRoutingGraphConfiguration.availableRoutes;

/**
 * @author JTB
 */
public class LlmRoutingNode{

    //大模型问题分类路由节点的返回结果
    public record LlmRoutingResponse(String reasoning, String selection) {}

    //大模型问题分类路由节点
    public static class LlmCallRouterNode implements NodeAction {

        private final ChatClient client;
        private final String inputTextKey;

        public LlmCallRouterNode(ChatClient client, String inputTextKey) {
            this.client = client;
            this.inputTextKey = inputTextKey;
        }

        @Override
        public Map<String, Object> apply(OverAllState state) {
            String inputText = (String) state.value(inputTextKey).orElse("");
            String selectorPrompt = String.format("""
          Analyze the input and select the most appropriate support team from these options: %s
          First explain your reasoning, then provide your selection in this JSON format:

          \\{
              "reasoning": "Brief explanation of why this ticket should be routed to a specific team.
                          Consider key terms, user intent, and urgency level.",
              "selection": "The chosen team name"
          \\}

          Input: %s""", availableRoutes.keySet(), inputText);

            LlmRoutingResponse llmRoutingResponse = client.prompt(selectorPrompt).call()
                    .entity(LlmRoutingResponse.class);
            Map<String, Object> selectionLlmMap = new HashMap<>();
            selectionLlmMap.put("selectionLlm", llmRoutingResponse.selection);
            return selectionLlmMap;
        }
    }
}
